import os
import math
import pickle
import numpy
from scipy import interpolate
from CalTorqueGEDIT import OpenCalTorqueXLS, GetSheetData
from lmfit import minimize, Parameters	
import matplotlib
from matplotlib import pyplot as PLT
													# Will then have to seperate into in and out
#_________________________________________________________________________ Get data
datapath = 'dataGEDIT/'
lsdir = os.listdir(datapath)

for f in lsdir:
   if (f.split('.')[-1] not in ['xlsx', 'xls']):
      continue
   print f
   book = OpenCalTorqueXLS(datapath+f)
   data = GetSheetData(book)


x = numpy.array(data['position'])
y = numpy.array(data['torque_inoz'])

#_________________________________________________________________________ Get pickled averages

def AverageTorqueInterp(x, offset = 0):			### if put x as argument, would there be an offset arg too? OFFSET FOR AVERAGE
	file = open('pickled_average_in.pkl', 'rb')
	avgdata = pickle.load(file)
	file.close()
	
	Xavg = numpy.array(avgdata['x'])
	Yavg = numpy.array(avgdata['y'])
	
	x2 = []
	f = interpolate.interp1d(Xavg, Yavg)
	if offset != 0:
		for i, j in enumerate(x):
			x2.append(x[i] - offset)
	else:
		x2 = x

	#Xin = []
	#for i in x2:
	#	if i > 5 and i < 395:
	#		Xin.append(i)
	#print len(Xin)
	return f(x2)


#_________________________________________________________________________ Get residual

def Residual(position, torque, offset = 0):
	trim_pos = []
	trim_data = []
	
	for i, x in enumerate(position):		## trims position data to ]5, 395[
		if x > 5 and x < 395:
			trim_pos.append(x)
		if x in trim_pos:					## trims torque data according to trimmed position
			trim_data.append(torque[i])
	
	avg_data = AverageTorqueInterp(trim_pos, offset)	## gets model data from trimmed position
	
	residual = 0
	for i, y in enumerate(trim_pos):
		residual += (trim_data[i] - avg_data[i])**2		## calc residual
	
	return residual**0.5


#_________________________________________________________________________ Plot residual vs offset

## creating offset values
offsetdata = numpy.linspace(-5, 5, 101)


## function that finds the residual for certain offsets
def ResfromOff(position, torque, offset):		### hm, takes a couple seconds to loop through this... is there a more time efficient way to execute this?
	resoff = []
	
	for i, k in enumerate(offset):
		res = Residual(position, torque, offset[i])
		resoff.append(res)
	
	return resoff

## plot residual vs offset
fig = PLT.figure(figsize = (15, 8), dpi = 150)
box = [0.14, 0.14, 0.76, 0.76]
ax1 = fig.add_axes(box)
ax1.set_ylabel('Residual (in_oz?)')
ax1.set_xlabel('Offset (cm)')
PLT.title('Residual versus Offset')
ax1.plot(offsetdata, ResfromOff(x, y, offsetdata), '-', color = 'b', linewidth = 2)
print 'Saving residual plot.'
PLT.savefig('ResidualPlot.png')


## min/max offsets, 1cm bins, trim average